"""
    create by IAmFiveHigh on 2023/8/20
"""
from typing import Union, Tuple, List

from .Vector import Vector
from LA.__globals import is_zero


class Matrix:
    def __init__(self, list2d):
        # 如果list2d是二维数组, 复制数组
        # 如果List2d是向量数组, 复制向量的底层列表组成数组
        if isinstance(list2d[0], list):
            self._values = [row[:] for row in list2d]
        elif isinstance(list2d[0], Vector):
            self._values = [row.under_list() for row in list2d]

    def __repr__(self):
        return "Matrix(\n{}\n)".format("\n".join(map(str, self._values)))

    __str__ = __repr__

    def item(self, r: int, c: int) -> Union[int, float]:
        # 第r行第c列的元素
        return self._values[r][c]

    def shape(self) -> Tuple[int, int]:
        return len(self._values), len(self._values[0])

    def row_num(self) -> int:
        # 行数
        return self.shape()[0]

    __len__ = row_num

    def col_num(self) -> int:
        # 列数
        return self.shape()[1]

    def size(self) -> int:
        r, c = self.shape()
        return r * c

    def row_vector(self, index: int) -> 'Vector':
        v = Vector(self._values[index])
        return v

    def col_vector(self, index: int) -> 'Vector':
        v = Vector([r[index] for r in self._values])
        return v

    def __add__(self, other: 'Matrix') -> 'Matrix':
        # 矩阵加法
        self.__guard_shape_same(other)
        return Matrix([
            [a + b for a, b in zip(self.row_vector(i), other.row_vector(i))]
            for i in range(self.row_num())
        ])

    def __sub__(self, other: 'Matrix') -> 'Matrix':
        # 矩阵减法
        self.__guard_shape_same(other)
        # 减法就是加上负的other
        return self + (-1 * other)

    def __mul__(self, k: Union[int, float]) -> 'Matrix':
        # 矩阵乘法
        return Matrix([
            [e * k for e in self.row_vector(i)]
            for i in range(self.row_num())
        ])

    def __rmul__(self, k: Union[int, float]) -> 'Matrix':
        return self * k

    def __truediv__(self, k: Union[int, float]) -> 'Matrix':
        # 矩阵除法
        if is_zero(k):
            raise ZeroDivisionError(f"Error: 矩阵{self} / {k} 除数不能为零")
        return self * (1 / k)

    def __neg__(self) -> 'Matrix':
        return self * -1

    def dot(self, other: Union['Vector', 'Matrix']) -> Union['Vector', 'Matrix']:
        if isinstance(other, Vector):
            # 矩阵向量乘法
            if self.col_num() != len(other):
                raise Exception(f"Error: 向量{other}的行数应和矩阵{self}列数相等")
            return Vector([self.row_vector(i).dot(other) for i in range(self.row_num())])
        elif isinstance(other, Matrix):
            # 矩阵矩阵乘法
            if self.col_num() != other.row_num():
                raise Exception(f"Error: 左矩阵{self}列数应和右矩阵{other}行数相等")
            return Matrix([[self.row_vector(i).dot(other.col_vector(j)) for j in range(other.col_num())]
                           for i in range(self.row_num())])

    # 判断矩阵是不是对角矩阵
    def is_diagonal(self):
        # 对角矩阵必须是方阵
        if self.col_num() != self.row_num():
            return False
        for i in range(self.row_num()):
            for j in range(self.col_num()):
                if i != j:
                    if self.item(i, j) != 0:
                        return False
        return True

    # 矩阵次方
    def __pow__(self, power, modulo=None):
        if self.is_diagonal():
            # 如果是对角矩阵，直接对对角线上值做次方运算
            return Matrix.diagonal([self.item(i, i) ** power for i in range(self.row_num())])
        else:
            # 否则循环多次矩阵乘法
            m = self.copy()
            for i in range(power - 1):
                m = m.dot(self.copy())
            return m

    # 返回一个和自己一样的矩阵
    def copy(self):
        return Matrix([
            [self.item(i, j) for j in range(self.col_num())]
            for i in range(self.row_num())
        ])

    # 矩阵的转置
    def T(self) -> 'Matrix':
        return Matrix([[e for e in self.col_vector(i)]
                       for i in range(self.col_num())])

    # 生成一个零矩阵
    @classmethod
    def zero(cls, r: int, c: int) -> 'Matrix':
        return cls([[0] * c for _ in range(r)])

    # 生成单位矩阵
    @classmethod
    def E(cls, n: int) -> 'Matrix':
        return cls([
            [1 if j == i else 0 for j in range(n)]
            for i in range(n)
        ])

    @classmethod
    def diagonal(cls, arr: List[Union[int, float]]):
        # 生成对角矩阵
        return cls([
            [arr[i] if j == i else 0 for j in range(len(arr))]
            for i in range(len(arr))
        ])

    def __guard_shape_same(self, other):
        if self.shape() != other.shape():
            raise Exception(f"Error: 矩阵{self}与矩阵{other}形状不相等")
